import os
import librosa
import numpy as np

RAW_DATA_PATH = '../data/raw/非病理性'  # 原始音频文件的路径
PROCESSED_DATA_PATH = '../data/raw/processed/'  # 预处理后音频文件保存的路径
TARGET_SAMPLE_RATE = 22050  # 目标采样率
CLIP_DURATION = 4.0  # 音频剪辑时长（秒）


def preprocess_audio(file_path):
    """
    预处理音频文件：加载音频、重采样、去除静音、剪辑和归一化。
    Args:
        file_path (str): 音频文件路径。
    Returns:
        y (numpy array): 预处理后的音频数据。
    """
    y, sr = librosa.load(file_path, sr=None)  # 加载音频文件
    y = librosa.resample(y, orig_sr=sr, target_sr=TARGET_SAMPLE_RATE)  # 重采样到22050Hz

    # 去除静音部分
    y, _ = librosa.effects.trim(y)

    y = librosa.util.normalize(y)  # 归一化音频数据

    # 剪辑音频数据
    if len(y) > TARGET_SAMPLE_RATE * CLIP_DURATION:
        y = y[:int(TARGET_SAMPLE_RATE * CLIP_DURATION)]
    else:
        y = np.pad(y, (0, max(0, int(TARGET_SAMPLE_RATE * CLIP_DURATION) - len(y))), 'constant')

    return y


def save_processed_audio(y, file_name):
    """
    保存预处理后的音频数据为.npy文件。
    Args:
        y (numpy array): 预处理后的音频数据。
        file_name (str): 保存的文件名。
    """
    np.save(os.path.join(PROCESSED_DATA_PATH, file_name), y)


def preprocess_all_audio():
    """
    遍历原始音频文件夹，预处理所有音频文件并保存。
    """
    if not os.path.exists(PROCESSED_DATA_PATH):
        os.makedirs(PROCESSED_DATA_PATH)

    label = "非病理性"

    for file in os.listdir("./data/raw/非病理性"):
        if file.endswith('.wav'):  # 处理所有.wav文件
            file_path = os.path.join("./data/raw/非病理性", file)
            y = preprocess_audio(file_path)  # 预处理音频
            save_processed_audio(y, f"{label}_{file.replace('.wav', '.npy')}")  # 保存预处理后的音频


if __name__ == '__main__':
    preprocess_all_audio()  # 执行预处理
